# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""模板参数定义

本模块定义了模板系统的核心数据结构，用于支持泛型类型和函数模板。
"""

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, TYPE_CHECKING
from enum import Enum

if TYPE_CHECKING:
    from mcpy.symbol_meta import TypeInfo
    from mcpy.ir import IRNode
    from mcpy.symbol_registry import SymbolRegistry


class ConstraintKind(Enum):
    """类型约束种类"""

    NONE = "none"  # 无约束
    UNION = "union"  # 联合类型约束：必须在允许的类型列表中
    PROTOCOL = "protocol"  # 协议约束：必须实现特定方法
    TEMPLATE = "template"  # 模板约束：必须匹配特定模板模式


@dataclass
class TemplateParam:
    """模板参数定义

    表示一个模板参数（如 T, K, V）及其约束信息。

    示例：
        # 无约束的通用参数
        T = TemplateParam(name="T")

        # 数值类型约束
        T_numeric = TemplateParam(
            name="T_numeric",
            constraint_kind=ConstraintKind.UNION,
            allowed_types=["int", "float"],
        )

        # 可迭代类型约束
        T_iter = TemplateParam(
            name="T_iter",
            constraint_kind=ConstraintKind.PROTOCOL,
            required_methods=["__iter__"],
        )

    Attributes:
        name: 参数名称（如 "T", "K", "V", "T_numeric"）
        constraint_kind: 约束类型
        default_type: 默认类型（Python 类型名，如 "Any"）
        default_cpp_type: 默认 C++ 类型（如 "mc::variant"）
        cpp_type_prefix: C++ 类型前缀（如 "typename"）
        allowed_types: union 约束的允许类型列表
        required_methods: protocol 约束的必需方法列表
        template_pattern: template 约束的模板模式
        description: 参数描述（用于文档和错误消息）
    """

    name: str
    constraint_kind: ConstraintKind = ConstraintKind.NONE
    is_variadic: bool = False  # 是否为可变参数（如 Args...）
    default_type: Optional[str] = None
    default_cpp_type: Optional[str] = None
    cpp_type_prefix: str = ""

    # 约束配置
    allowed_types: List[str] = field(default_factory=list)
    required_methods: List[str] = field(default_factory=list)
    template_pattern: Optional[str] = None
    bound: Optional[str] = None  # Protocol bound 名称（如 "Integer", "Iterable"）

    # 元数据
    description: Optional[str] = None

    # 延迟解析的 bound 信息（缓存）
    _resolved_bound: Optional[Dict[str, Any]] = field(
        default=None, repr=False, compare=False
    )

    # 参数传递策略（用于 C++ 代码生成）
    # "value": 按值传递（T x）
    # "const_ref": 按常量引用传递（const T& x）
    # "rvalue_ref": 按右值引用传递（T&& x，用于完美转发）
    # "auto": 自动选择（根据类型大小决定）
    pass_by: str = "auto"

    def __post_init__(self):
        """验证配置一致性"""
        if self.constraint_kind == ConstraintKind.UNION and not self.allowed_types:
            raise ValueError(f"Union constraint '{self.name}' must have allowed_types")
        if (
            self.constraint_kind == ConstraintKind.PROTOCOL
            and not self.required_methods
        ):
            raise ValueError(
                f"Protocol constraint '{self.name}' must have required_methods"
            )
        if (
            self.constraint_kind == ConstraintKind.TEMPLATE
            and not self.template_pattern
        ):
            raise ValueError(
                f"Template constraint '{self.name}' must have template_pattern"
            )

    def has_constraint(self) -> bool:
        """是否有类型约束"""
        return self.constraint_kind != ConstraintKind.NONE

    def matches_type(
        self,
        type_info: "TypeInfo",
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> bool:
        """检查类型是否满足约束

        Args:
            type_info: 要检查的类型信息
            node: 上下文 IR 节点，用于作用域查找（可选）
            registry: 符号注册表，用于全局类型查找（可选）

        Returns:
            bool: 类型是否满足约束

        Notes:
            - node 和 registry 用于查找 Protocol 定义和类型的方法
        """
        # 如果有 bound，需要延迟解析（避免循环依赖）
        if self.bound and not self._resolved_bound:
            self._resolve_bound_lazy(node, registry)

        if self.constraint_kind == ConstraintKind.NONE:
            # 无约束，任何类型都匹配
            return True

        elif self.constraint_kind == ConstraintKind.UNION:
            # 联合约束：类型必须在允许列表中
            return type_info.name in self.allowed_types

        elif self.constraint_kind == ConstraintKind.PROTOCOL:
            # 协议约束：类型必须实现所需方法
            return self._check_protocol_constraint(type_info, node, registry)

        elif self.constraint_kind == ConstraintKind.TEMPLATE:
            # 模板约束：类型必须匹配模板模式
            return self._check_template_constraint(type_info)

        return False

    def _resolve_bound_lazy(
        self,
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> None:
        """延迟解析 bound（避免在初始化时循环依赖）

        Args:
            node: 上下文 IR 节点
            registry: 符号注册表
        """
        if not self.bound or self._resolved_bound is not None:
            return

        # 调用静态方法解析（传递上下文）
        bound_info = self._resolve_protocol_bound(self.bound, node, registry)
        if bound_info:
            # 更新约束信息
            self.constraint_kind = bound_info["constraint_kind"]
            if "required_methods" in bound_info:
                self.required_methods = bound_info["required_methods"]
            if "allowed_types" in bound_info:
                self.allowed_types = bound_info["allowed_types"]
            self._resolved_bound = bound_info
        else:
            # 解析失败，标记为已解析（避免重复尝试）
            self._resolved_bound = {}

    def _check_protocol_constraint(
        self,
        type_info: "TypeInfo",
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> bool:
        """检查协议约束

        从符号表查询类型定义，获取其方法集合，验证是否满足 Protocol 要求。

        Args:
            type_info: 要检查的类型信息
            node: 上下文 IR 节点
            registry: 符号注册表

        Returns:
            bool: 类型是否满足协议约束
        """
        if type_info.methods:
            type_methods = set(type_info.methods.keys())
            required_methods_set = set(self.required_methods)
            return required_methods_set.issubset(type_methods)

        # 从符号表查询类型的方法集合
        type_methods = self._get_type_methods_from_registry(
            type_info.name, node, registry
        )

        required_methods_set = set(self.required_methods)
        return required_methods_set.issubset(type_methods)

    def _check_template_constraint(self, type_info: "TypeInfo") -> bool:
        """检查模板约束"""
        from mcpy.symbol_meta import TypeKind

        # 简单的模板模式匹配
        if self.template_pattern and self.template_pattern.startswith("Pointer["):
            return type_info.kind == TypeKind.GENERIC and type_info.name == "Pointer"

        return False

    @staticmethod
    def _get_type_methods_from_registry(
        type_name: str,
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> set:
        """从符号表查询类型的方法集合

        Args:
            type_name: 类型名称（如 "int", "list", "MyClass"）
            node: 上下文 IR 节点（用于作用域查找）
            registry: 符号注册表（用于全局查找）

        Returns:
            set: 类型支持的方法名集合

        Notes:
            查找顺序：
            1. 从 node 的作用域查找（局部类型）
            2. 从 registry 查找（模块类型）
            3. 从全局符号表查找（BUILTIN 类型）
        """
        from mcpy.symbol_meta import SymbolKind

        type_symbol = None

        # 1. 从 node 的作用域查找
        if node is not None:
            try:
                type_symbol = node.lookup_symbol(
                    type_name, symbol_kinds=[SymbolKind.CLASS]
                )
            except Exception:
                pass

        # 2. 从 registry 查找
        if type_symbol is None and registry is not None:
            type_symbol = registry.lookup_type_symbol(type_name)

        # 3. 从全局符号表查找
        if type_symbol is None:
            try:
                from mcpy.symbol_registry import (
                    get_global_registry,
                    is_global_registry_initializing,
                )

                if not is_global_registry_initializing():
                    global_registry = get_global_registry()
                    if registry is not global_registry:
                        type_symbol = global_registry.lookup_type_symbol(type_name)
            except Exception:
                pass

        # 提取方法集合
        all_methods = set()

        if type_symbol:
            # 1. 从 type_info.methods 获取（用户自定义类型）
            if type_symbol.type_info and type_symbol.type_info.methods:
                all_methods.update(type_symbol.type_info.methods.keys())

            # 2. 从 type_config 获取方法和 Protocol 信息
            if hasattr(type_symbol, "type_config") and type_symbol.type_config:
                # 2.1 直接定义的方法
                config_methods = type_symbol.type_config.get("methods", {})
                if config_methods:
                    all_methods.update(config_methods.keys())

                # 2.2 从 satisfies_protocols 推导方法
                # 对于原生 C++ 类型（int, float, bool）或模板类型（list, dict），
                # 如果配置中声明了满足的 Protocol，从 Protocol 定义推导方法
                satisfies = type_symbol.type_config.get("satisfies_protocols", [])
                if satisfies:
                    for protocol_name in satisfies:
                        protocol_methods = TemplateParam._get_protocol_methods(
                            protocol_name, node, registry
                        )
                        all_methods.update(protocol_methods)

        return all_methods

    @staticmethod
    def _get_protocol_methods(
        protocol_name: str,
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> set:
        """从 Protocol 定义获取 required_methods

        Args:
            protocol_name: Protocol 名称
            node: 上下文 IR 节点
            registry: 符号注册表

        Returns:
            Protocol 要求的方法名集合
        """
        from mcpy.symbol_meta import SymbolKind

        protocol_symbol = None

        # 查找 Protocol（与 _resolve_protocol_bound 类似的逻辑）
        if node is not None:
            try:
                protocol_symbol = node.lookup_symbol(
                    protocol_name, symbol_kinds=[SymbolKind.CLASS]
                )
            except Exception:
                pass

        if protocol_symbol is None and registry is not None:
            protocol_symbol = registry._primitive_types.get(
                protocol_name
            ) or registry._template_types.get(protocol_name)

        if protocol_symbol is None:
            try:
                from mcpy.symbol_registry import (
                    get_global_registry,
                    is_global_registry_initializing,
                )

                if not is_global_registry_initializing():
                    global_registry = get_global_registry()
                    protocol_symbol = global_registry._primitive_types.get(
                        protocol_name
                    ) or global_registry._template_types.get(protocol_name)
            except Exception:
                pass

        # 提取 required_methods
        if protocol_symbol and protocol_symbol.type_info:
            if protocol_symbol.type_info.is_protocol:
                return set(protocol_symbol.type_info.required_methods)

        return set()

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典（用于序列化）"""
        result = {
            "name": self.name,
            "constraint_kind": self.constraint_kind.value,
        }

        if self.default_type:
            result["default_type"] = self.default_type
        if self.default_cpp_type:
            result["default_cpp_type"] = self.default_cpp_type
        if self.cpp_type_prefix:
            result["cpp_type_prefix"] = self.cpp_type_prefix
        if self.allowed_types:
            result["allowed_types"] = self.allowed_types
        if self.required_methods:
            result["required_methods"] = self.required_methods
        if self.template_pattern:
            result["template_pattern"] = self.template_pattern
        if self.bound:
            result["bound"] = self.bound
        if self.description:
            result["description"] = self.description
        if self.pass_by != "auto":
            result["pass_by"] = self.pass_by

        return result

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TemplateParam":
        """从字典创建（用于反序列化）"""
        constraint_kind_str = data.get("constraint_kind", "none")
        constraint_kind = ConstraintKind(constraint_kind_str)

        return cls(
            name=data["name"],
            constraint_kind=constraint_kind,
            default_type=data.get("default_type"),
            default_cpp_type=data.get("default_cpp_type"),
            cpp_type_prefix=data.get("cpp_type_prefix", ""),
            allowed_types=data.get("allowed_types", []),
            required_methods=data.get("required_methods", []),
            template_pattern=data.get("template_pattern"),
            bound=data.get("bound"),
            description=data.get("description"),
            pass_by=data.get("pass_by", "auto"),
        )

    @classmethod
    def from_type_var_config(cls, name: str, config: Dict[str, Any]) -> "TemplateParam":
        """从 TYPE_VARS 配置创建模板参数

        Args:
            name: 参数名称
            config: TYPE_VARS 中的配置字典

        Returns:
            TemplateParam 实例

        示例：
            # 方式1：具体类型约束
            config = {
                "description": "数值类型",
                "constraints": ["int", "float"],
            }
            param = TemplateParam.from_type_var_config("T_numeric", config)

            # 方式2：引用 Protocol（推荐）
            config = {
                "description": "整数类型",
                "bound": "Integer",
            }
            param = TemplateParam.from_type_var_config("T_int", config)
        """
        # 获取配置项
        bound = config.get("bound")  # Protocol 引用
        constraints = config.get("constraints", [])
        constraint_kind_str = config.get("constraint_kind", "")
        required_methods = config.get("required_methods", [])
        template_pattern = config.get("template_pattern")

        # 推断约束类型
        # 注意：如果有 bound，约束类型将在首次使用时延迟解析（避免循环依赖）
        if bound:
            # 有 bound 时，先设为 NONE，等到 matches_type 时再延迟解析
            constraint_kind = ConstraintKind.NONE
        elif constraint_kind_str:
            constraint_kind = ConstraintKind(constraint_kind_str)
        elif constraints:
            constraint_kind = ConstraintKind.UNION
        elif required_methods:
            constraint_kind = ConstraintKind.PROTOCOL
        elif template_pattern:
            constraint_kind = ConstraintKind.TEMPLATE
        else:
            constraint_kind = ConstraintKind.NONE

        return cls(
            name=name,
            constraint_kind=constraint_kind,
            default_type=config.get("default_type"),
            default_cpp_type=config.get("default_cpp_type"),
            cpp_type_prefix=config.get("cpp_type_prefix", ""),
            allowed_types=constraints,
            required_methods=required_methods,
            template_pattern=template_pattern,
            bound=bound,  # 保存 bound 名称，延迟解析
            description=config.get("description"),
            pass_by=config.get("pass_by", "auto"),
        )

    @staticmethod
    def _resolve_protocol_bound(
        protocol_name: str,
        node: Optional["IRNode"] = None,
        registry: Optional["SymbolRegistry"] = None,
    ) -> Optional[Dict[str, Any]]:
        """解析 Protocol 引用，从符号表获取协议定义

        Args:
            protocol_name: Protocol 名称（如 "Integer", "Comparable", "Iterable"）
            node: 上下文 IR 节点
            registry: 符号注册表

        Returns:
            Protocol 配置字典，如果找不到返回 None

        Notes:
            - 支持泛型协议，如 "Iterable" 会匹配 "Iterable[T]"
            - 从符号表查询，支持 builtin 和 stub 定义的 Protocol
            - 查找顺序：node 作用域 → registry → 全局符号表
        """
        from mcpy.symbol_meta import SymbolKind

        protocol_symbol = None

        # 1. 从 node 的作用域查找
        if node is not None:
            try:
                protocol_symbol = node.lookup_symbol(
                    protocol_name, symbol_kinds=[SymbolKind.CLASS]
                )
            except Exception:
                pass

        # 2. 从 registry 查找
        if protocol_symbol is None and registry is not None:
            protocol_symbol = registry._primitive_types.get(
                protocol_name
            ) or registry._template_types.get(protocol_name)

        # 3. 从全局符号表查找
        if protocol_symbol is None:
            try:
                from mcpy.symbol_registry import (
                    get_global_registry,
                    is_global_registry_initializing,
                )

                if not is_global_registry_initializing():
                    global_registry = get_global_registry()
                    protocol_symbol = global_registry._primitive_types.get(
                        protocol_name
                    ) or global_registry._template_types.get(protocol_name)
            except Exception:
                pass

        # 如果找不到，可能是泛型 Protocol，尝试常见模式
        if not protocol_symbol:
            # 尝试从全局符号表查找泛型模式
            try:
                from mcpy.symbol_registry import get_global_registry

                global_registry = get_global_registry()
                for pattern in [f"{protocol_name}[T]", f"{protocol_name}[K, V]"]:
                    protocol_symbol = global_registry._template_types.get(pattern)
                    if protocol_symbol:
                        break
            except Exception:
                pass

        if not protocol_symbol or not protocol_symbol.type_info:
            return None

        # 检查是否是 Protocol
        if not protocol_symbol.type_info.is_protocol:
            return None

        # 提取约束信息
        required_methods = protocol_symbol.type_info.required_methods

        # 判断约束类型
        if required_methods:
            return {
                "constraint_kind": ConstraintKind.PROTOCOL,
                "required_methods": required_methods,
            }
        else:
            # 某些 Protocol 可能定义为 Union 类型（从 type_config 获取）
            if hasattr(protocol_symbol, "type_config") and protocol_symbol.type_config:
                allowed_types = protocol_symbol.type_config.get("allowed_types", [])
                if allowed_types:
                    return {
                        "constraint_kind": ConstraintKind.UNION,
                        "allowed_types": allowed_types,
                    }

        return None

    def __repr__(self) -> str:
        """可读的字符串表示"""
        if self.constraint_kind == ConstraintKind.NONE:
            return f"TemplateParam({self.name})"
        elif self.constraint_kind == ConstraintKind.UNION:
            types_str = ", ".join(self.allowed_types)
            return f"TemplateParam({self.name}: {types_str})"
        elif self.constraint_kind == ConstraintKind.PROTOCOL:
            methods_str = ", ".join(self.required_methods)
            return f"TemplateParam({self.name}: protocol[{methods_str}])"
        else:
            return f"TemplateParam({self.name}: {self.constraint_kind.value})"
